import argparse
import json
import math
import os

import cv2
import numpy as np
from numpy import matrix
from scipy import ndimage


from tqdm import tqdm


#list_xml = os.listdir(path)
from xml.dom.minidom import parse
import xml.dom.minidom
import xml.etree.ElementTree as ET

pi = 3.14

parser = argparse.ArgumentParser('rec')
parser.add_argument("--xmlpath", type=str, default='D:/BaiduNetdiskDownload/vin_6.28_roLabel/vin6.28/simple/xml')
parser.add_argument("--imgpath", type=str, default='D:/BaiduNetdiskDownload/vin_6.28_roLabel/vin6.28/simple/1')
parser.add_argument("--savepath", type=str, default='D:/test')

args = parser.parse_args()


def rotate(ps, m):
    pts = np.float32(ps).reshape([-1, 2])  # 要映射的点
    pts = np.hstack([pts, np.ones([len(pts), 1])]).T
    target_point = np.dot(m, pts)
    target_point = [[int(target_point[0][x]),int(target_point[1][x])] for x in range(len(target_point[0]))]
    return target_point

def rotate_img_and_point(img,points,angle,center_x,center_y,resize_rate=1.0):
    h,w,c = img.shape
    M = cv2.getRotationMatrix2D((center_x,center_y), angle, resize_rate)
    res_img = cv2.warpAffine(img, M, (w, h))
    out_points = rotate(points, M)
    return res_img,out_points

def rotate_point(points,angle,center_x,center_y,resize_rate=1.0):
    M = cv2.getRotationMatrix2D((center_x, center_y), angle, resize_rate)
    out_points = rotate(points, M)
    return  out_points


def get_rec_img(xml_path,img_path,savepath):

    list_img = os.listdir(img_path)
    for i in tqdm(range(len(list_img))):
        xml_name = list_img[i].split(".")[0] + ".xml"
        if os.path.exists(os.path.join(xml_path, xml_name)):
            # 打开xml文档
            # print(list_xml[i])
            box = {"x1": 0, "y1": 0, "x2": 0, "y2": 0, "angle": 0}
            dom = xml.dom.minidom.parse(os.path.join(xml_path, xml_name))

            # 得到文档元素对象
            root = dom.documentElement
            if len(root.getElementsByTagName('object')) > 0:
                object = root.getElementsByTagName('object')[0]
                if len(object.getElementsByTagName('robndbox')) == 0:
                    continue
                robndbox = object.getElementsByTagName('robndbox')[0]

                img_weight = int(
                    root.getElementsByTagName('size')[0].getElementsByTagName('width')[0].childNodes[0].data)
                img_height = int(
                    root.getElementsByTagName('size')[0].getElementsByTagName('height')[0].childNodes[0].data)
                # print(img_height,img_weight)
                cx = robndbox.getElementsByTagName('cx')[0]
                x0 = int(float(cx.childNodes[0].data))
                cy = robndbox.getElementsByTagName('cy')[0]
                y0 = int(float(cy.childNodes[0].data))
                w = int(float(robndbox.getElementsByTagName('w')[0].childNodes[0].data))
                h = int(float(robndbox.getElementsByTagName('h')[0].childNodes[0].data))
                # print(x0, y0, w, h)
                ori = {"x1": x0 - int(w / 2.0), "y1": y0 - int(h / 2.0), "x2": x0 + int(w / 2.0),
                       "y2": y0 + int(h / 2.0)}

                box['angle'] = float(robndbox.getElementsByTagName('angle')[0].childNodes[0].data)
                img = cv2.imread(os.path.join(img_path,list_img[i]))

                if w < h:
                    # image = ndimage.rotate(img, ((box['angle']) / math.pi * 180)+90.0)
                    image, p = rotate_img_and_point(img, points=[x0, y0], angle=((box['angle']) / math.pi * 180) + 90.0,
                                                    center_x=int(img_weight / 2.0), center_y=int(img_height / 2.0),
                                                    resize_rate=1.0)
                    # print("p",p)
                    temp = w
                    w = h
                    h = temp


                else:

                    image, p = rotate_img_and_point(img, points=[x0, y0], angle=((box['angle']) / math.pi * 180),
                                                    center_x=int(img_weight / 2.0), center_y=int(img_height / 2.0),
                                                    resize_rate=1.0)


                lv = [[1, 1]]
                res_sp = image.shape
                for num in lv:

                    result_p = []
                    if int((p[0][1] - int(h / 2.0)) * num[0]) < 0:
                        result_p.append(0)
                    else:
                        result_p.append(int((p[0][1] - int(h / 2.0)) * num[0]))

                    if int((p[0][1] + int(h / 2.0)) * num[1]) > res_sp[0]:
                        result_p.append(res_sp[0])
                    else:
                        result_p.append(int((p[0][1] + int(h / 2.0)) * num[1]))

                    if int((p[0][0] - int(w / 2.0)) * num[0]) < 0:
                        result_p.append(0)
                    else:
                        result_p.append(int((p[0][0] - int(w / 2.0)) * num[0]))

                    if int((p[0][0] + int(w / 2.0)) * num[1]) > res_sp[1]:
                        result_p.append(res_sp[1])
                    else:
                        result_p.append(int((p[0][0] + int(w / 2.0)) * num[1]))

                    img = image[result_p[0]:result_p[1], result_p[2]:result_p[3]]
                    if img.size != 0:
                        cv2.imwrite(os.path.join(savepath,list_img[i]), img)




get_rec_img(args.xmlpath,args.imgpath,args.savepath)